function [output cache] = functions_matrix_completion_polar(fname, model, searchDir, data, params, cache)
% Set of functions that are specific to matrix completion problem
%
% Total space = St(n,p)\times Cone(p)\times St(m,p)
% St(n,p): Stiefel manifold
% Cone(p): Positive definite matrices of size p\times p
%
% Usages:
%
% Cost function:                            f
% Gradient on the total space:              grad 
% Riemmanian Hessian on the total space:    hess 
% Dual variable:                            dual_var
% model is the current iterate
% eta is a search direction
% params is a cell array that contains {I,J,trueDists,EIJ} (see lowrank_dist_completion.m)
% cache is a cell array to reduce redundant computations
%
% Authors:
% Bamdev Mishra and Gilles Meyer
% {b.mishra,g.meyer}@ulg.ac.be

switch lower(fname),   
  
    case 'predict', % Compute predictions
        
        output = partXY((model.U*model.B)', model.V', data.rows, data.cols, length(data.rows))'; % Computationally expensive
        
        
    case 'f', % Objective function
        
        preds = functions_matrix_completion_polar('predict', model, [], data);
        errors = (preds - data.entries);
        
        % Squared Error + trace norm
        output = (errors'*errors) + model.lambda*trace(model.B);
        
        
        
    case 'grad_f', % Compute gradient
        
        S = functions_matrix_completion_polar('dual_var', model, [], data);
        SU = S'*model.U; % Computationally expensive
        SV = S*model.V; % Computationally expensive
        USV = model.U'*SV;
        grad.U = SV*model.B;
        grad.U = proj_tangent_space_stiefel(grad.U,model.U);
        grad.B = model.B*(symm(USV)) * model.B + model.lambda*model.B*model.B;
        grad.V = SU*model.B;
        grad.V = proj_tangent_space_stiefel(grad.V,model.V);
        
        % Cache the values of S and gradient
        % for subsequent computation of the Hessian
        cache.S = S;
        cache.grad = grad;
        cache.SU = SU;
        cache.SV = SV;
        
        output = grad;
        
    case 'hessian', % Compute hessian
        
        if isfield(cache,'S'),
            S = cache.S;
        else
            S = functions_matrix_completion_polar('dual_var', model, [], data);
        end
        
        if isfield(cache,'grad'),
            grad = cache.grad;
        else
            grad = functions_matrix_completion_polar('grad_f', model, [], data);
        end
        
        if isfield(cache, 'SU')
            SU = cache.SU;
            SV = cache.SV;
        else
            SU = S'*model.U;
            SV = S*model.V;
        end
        
        s_ubv = partXY([(searchDir.U*model.B + model.U*searchDir.B)  model.U*model.B ]', [model.V   searchDir.V]', data.rows, data.cols, length(data.cols))'; % Computationally expensive
        if isfield(cache,'S_star')
            updateSparse(cache.S_star, 2 * s_ubv);
            S_star = cache.S_star;
        elseif isfield(model,'sparse_structure2') % If you have another sparse structure, why not use it!
            updateSparse(model.sparse_structure2, 2 * s_ubv);
            S_star = model.sparse_structure2;
        else
            warning('WarnTests:convertTest',...
                'Forming a sparse matrix naively.\nFor efficiency supply a sparse skeleton, "model.sparse_structure2" in the input.\n');
            S_star = sparse(data.rows, data.cols, 2 * s_ubv, model.d1, model.d2);
            cache.S_star = S_star;
        end
        
        
        
        SDirU = S'*searchDir.U; % Computationally expensive
        SDirV = S * searchDir.V; % Computationally expensive
        S_starU = S_star'*model.U; % Computationally expensive
        S_starV = S_star*model.V; % Computationally expensive
        USV = model.U'*SV;
        
        hessian.U = SV*searchDir.B + SDirV*model.B + S_starV*model.B;
        % Projection onto the tangent space to the Stiefel Manifold
        hessian.U = proj_tangent_space_stiefel(hessian.U, model.U) - ...
            proj_tangent_space_stiefel(searchDir.U*symm(USV*model.B), model.U);
        
        symmUSV = symm(USV);
        
        
        hessian.B = searchDir.B*symmUSV*model.B  + ...
            model.B*symmUSV*searchDir.B  + ...
            model.B*symm(searchDir.U'*SV + model.U'*S_starV + model.U'*SDirV)*model.B + ...
            model.lambda*model.B*searchDir.B + model.lambda*searchDir.B*model.B;
        
        % Correction for non-constant metric
        hessian.B = hessian.B - symm(searchDir.B*(model.B\grad.B));
        
        hessian.V = SU*searchDir.B + SDirU*model.B + S_starU*model.B;
        % Projection onto the tangent space to the Stiefel Manifold
        hessian.V = proj_tangent_space_stiefel(hessian.V, model.V) - ...
            proj_tangent_space_stiefel(searchDir.V*symm(USV'*model.B),model.V);
        
        output = hessian;
        
    case 'dual_var', % Dual variable used to compute restart direction
        
        preds = functions_matrix_completion_polar('predict', model, [], data);
        errors = (preds - data.entries);
        if isfield(model,'sparse_structure')
            sparse_structure = model.sparse_structure; % We require a sparse skeleton
            updateSparse(sparse_structure, 2 * errors);
        else
            warning('WarnTests:convertTest',...
                'Forming a sparse matrix naively.\nFor efficiency supply a sparse skeleton, "model.sparse_structure" in the input.\n');
            sparse_structure = sparse(data.rows, data.cols, 2 *errors, model.d1, model.d2); % Computationally expensive
        end
        
        output = sparse_structure;
        
    otherwise,
        error('Unknown function %s');
end

